import random
import ConfigSpace as CS
import ConfigSpace.hyperparameters as CSH
import itertools

class Arg:
    cs_attributes = CS.ConfigurationSpace()
    #attribute_list = ['axis', 'keepdims', 'select_last_index']
    attribute_list = ['axis', 'keepdims']
    all_attributes = []

    def __init__(self):
        print('Arg Init--------------------------')
        self.init_all_attributes()
        self.init_cs()

    def init_all_attributes(self):
        for i in range(1, len(Arg.attribute_list)+1) :
            #此处取所有组合的可能值，存在列表里
            for attribute in itertools.combinations(Arg.attribute_list, i):
                l = list(attribute)
                Arg.all_attributes.append(l)

        #print('all_attributes: ', all_attributes)
        print('all attributes:')
        for a in Arg.all_attributes:
            print(a)

    def init_cs(self):
        length = len(Arg.all_attributes)
        attributes = CSH.UniformIntegerHyperparameter(name='combination', lower=0, upper=2*length)
        Arg.cs_attributes.add_hyperparameter(attributes)


    def get_random_combination(self):
        s = Arg.cs_attributes.sample_configuration()
        index = s.get_dictionary()['combination']
        print('get combination: ', index)

        length = len(Arg.all_attributes)

        #保证不带任何属性的case可以被覆盖到
        if index >= length:
            print('test no attributes')
            return []

        return Arg.all_attributes[index]       

    def get_max_attributes_combination(self):
        return len(Arg.all_attributes)

    def get_attributes_by_index_or_random(self, shapes, index):

        print('get_attributes_by_index_or_random, index: ', index)    
            
        attributes = {}
        shape_out = []

        print('arg, shapes:', shapes[0])

        axis = 0
        keepdims = 0
        select_last_index = 0

        if index != -1:
            attr_combination = Arg.all_attributes[index]
        else:
            attr_combination = self.get_random_combination()

        if len(attr_combination) == 0:
            for i in range(len(shapes[0])):
                if i != 0:
                    shape_out.append(shapes[0][i])

            print('arg, shape_out:', shape_out)

            return attributes, shape_out    

        if 'axis' in attr_combination:
            axis = random.randint(-1*len(shapes[0]), len(shapes[0]) - 1)
            attributes['axis'] = axis

        if 'keepdims' in attr_combination:
            keepdims = random.randint(0, 1)
            attributes['keepdims'] = keepdims

        if 'select_last_index' in attr_combination:
            select_last_index = random.randint(0, 1)
            attributes['select_last_index'] = select_last_index    

        if axis < 0 :
            axis = len(shapes[0]) + axis

        for i in range(len(shapes[0])):
            if i != axis:
                shape_out.append(shapes[0][i])
                print('xxx shape_out:', shape_out)
            elif keepdims == 1:
                shape_out.append(1)
                print('yyy shape_out:', shape_out)    

        if len(shape_out) == 0:
            shape_out = [1]

        print('------ attributes:', attributes, 'shape_out:', shape_out)

        return attributes, shape_out 

    def input_shape_correct(self, shapes):
        return True

    def get_output_type(self, index):
        return 'tensor(int64)'        